package cz.cuni.amis.episodic.lisp.netcreators;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.io.PrintStream;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import jsmile.utils.CsvToSmileData;
import lombok.Cleanup;

import org.apache.commons.lang.StringUtils;

import smile.Network;
import cz.cuni.amis.episodic.lisp.LispPlanToCsvData;
import cz.cuni.amis.episodic.lisp.LispPlanToCsvData.LispToCsvResult;
import cz.cuni.amis.episodic.lisp.behan.LispTree;
import cz.cuni.amis.episodic.lisp.behan.TraceToEventVisitor;
import cz.cuni.amis.episodic.lisp.visitor.HHMMmpTraceVisitor;

/**
 * Sequence of altering Goal and Hidden nodes. Hidden nodes are learned by EM.
 * Finish nodes are needed to achieve power of hierarchical FSM.
 *
 * @author ik
 */
public class HHMMMurphyPaskin2001Creator extends NetCreator {

    int hiddenNodeStatesNum;

    public HHMMMurphyPaskin2001Creator(int hiddenNodeStatesNum) {
        this.hiddenNodeStatesNum = hiddenNodeStatesNum;
        filenameSuffix = "_hhmm_mp";
    }

    public List<List<String>> lispToTrainingData(File inputLispPlan, File smileTrainingData) throws IOException {
        // get filename without path and extension

        // read lisp to csv
        ByteArrayOutputStream csvOs = new ByteArrayOutputStream();
        LispPlanToCsvData.LispToCsvResult res = lispToCsv(inputLispPlan.getAbsolutePath(), csvOs, batchSize);
        csvOs.close();
        
        // write also the csv
        String smilePath = smileTrainingData.getAbsolutePath();
        File csvFile = new File(smilePath.substring(0, smilePath.lastIndexOf('.') + 1) + "csv");
        @Cleanup
        FileOutputStream fos = new FileOutputStream(csvFile);
        fos.write(csvOs.toByteArray());
        
        // transform csv to smile data format
        CsvToSmileData.convert(new InputStreamReader(new ByteArrayInputStream(csvOs.toByteArray())),
                smileTrainingData.getAbsolutePath());       

        return res.levelActions;
    }

    /**
     * UGLY cut and paste and edit from LispPlanToCsvData ... should be refactored.
     * @param inputLispPlan
     * @param outputPlanCsv
     * @param batchSize
     * @param translator
     * @return
     * @throws FileNotFoundException 
     */
    public static LispPlanToCsvData.LispToCsvResult lispToCsv(String inputLispPlan,
            OutputStream outputPlanCsv, int batchSize) throws FileNotFoundException {
        LispTree parser = new LispTree(inputLispPlan);
        LispPlanToCsvData.LispToCsvResult res = new LispPlanToCsvData.LispToCsvResult();

        HHMMmpTraceVisitor visitor = new HHMMmpTraceVisitor();
        parser.walkInForest(visitor);

        // find max depth
        int max = Integer.MIN_VALUE;
        for (String[] trace : visitor.getTraces()) {
            max = Math.max(max, trace.length);
        }

        @Cleanup
        PrintStream out = new PrintStream(outputPlanCsv);
        String delim = "\t";

        // each set contains actions found on the corresponding level in
        // hierarchy
        Set<String>[] levels = new HashSet[max];
        for (int i = 0; i < max; i++) {
            levels[i] = new HashSet<String>();
        }

        // print header
        String[] nodesInHeader = new String[max * 2 - 1];
        for (int i = 0; i < nodesInHeader.length - 1; i+=2) {
            int layerIx = max - (i / 2 + 1);
            nodesInHeader[i] = "G" + layerIx;
            nodesInHeader[i+1] = "F" + layerIx;
        }
        nodesInHeader[nodesInHeader.length - 1] = "O";
        List<String> header = new ArrayList<>(Arrays.asList("Batch", "Id"));
        header.addAll(Arrays.asList(nodesInHeader));
        out.println(StringUtils.join(header, delim));

        // compute finish nodes
 /*       List<List<String>> traces = new ArrayList<>();
        for (String[] trace : visitor.getTraces()) {
            traces.add(Arrays.asList(trace).subList(0, trace.length - 1));
        }
        boolean[][] finish = transformToHHMMFinish(traces);
   */     
        
        // output with padding

        int batchNum = 1;
        int linesInBatch = 0;
        String[] line = new String[max];
        int traceIx = 0;
        for (String[] trace : visitor.getTraces()) {
            linesInBatch++;
            out.print(batchNum + delim + Integer.toString(res.lineNum++)
                    + delim);
            
            // join goal nodes and finish nodes
            String[] outLine = new String[trace.length * 2 - 1];
            for(int i = 0; i < trace.length -1; i++) {
                outLine[2*i] = trace[i];
                outLine[2*i + 1] = visitor.getFinish().get(traceIx)[i] ? "finish" : "continue";
            }
            outLine[outLine.length-1] = trace[trace.length -1];
            
            out.println(StringUtils.join(outLine, delim));
           

            if (batchSize > 0 && linesInBatch == batchSize) {
                batchNum++;
                linesInBatch = 0;
            }

            // remember actions/goals found in each level
            for (int i = 0; i < trace.length; i++) {
                levels[i].add(trace[i]);
            }
            traceIx++;
        }

        // create result
        res.batchSize = batchSize;
        res.inputLispPlan = inputLispPlan;

        // sort actions by alphabet
        for (int i = 0; i < levels.length; i++) {
            List<String> actions = new ArrayList<>(levels[i]);
            Collections.sort(actions);
            res.levelActions.add(actions);
        }
        return res;
    }

      /**
     * Add finish states.
     * @return 
     */
    static boolean[][] transformToHHMMFinish(List<List<String>> data) {
         boolean[][] finish = new boolean[data.size()][];
        List<String> lastTrace = null;
        int traceIx = 0;
        for(List<String> trace : data) {
            if(lastTrace != null) {
                boolean[] finishData = new boolean[trace.size()];
                // compare
                for(int i = 0; i < trace.size(); i++) {
                    finishData[i] = !lastTrace.get(i).equals(trace.get(i));
                }
                finish[traceIx] = finishData;
                traceIx++;
            } else {
                // first trace, nothing to compare
            }
            lastTrace = trace;
        }
        // on last slice every state finishes
        boolean[] lastFinish = new boolean[lastTrace.size()];
        Arrays.fill(lastFinish, true);
        finish[finish.length-1] = lastFinish;
        return finish;
    }
    
    @Override
    public Network createNet(List<List<String>> levelActions) {
        List<String> statesF = Arrays.asList("continue", "finish");

        final int max = levelActions.size();

        Network net = new Network();
        String[] nodes = createNodeNames(max);
        for (int i = 0; i < max - 1; i++) {
            int id = max - i - 1;
            int nodeIx = 2 * i;
            String finishNode = "F" + (i + 1);

            // F ... finish node
            addNode(net, finishNode, statesF);
            deleteDefaultOutcomes(net, finishNode);

            // G
            addNode(net, nodes[nodeIx], levelActions.get(i));

            // H
            addNode(net, nodes[nodeIx + 1], hiddenNodeStatesNum);
            if (i > 0) {
                // H_(i-1) -> G_i
                net.addArc(nodes[nodeIx - 1], nodes[nodeIx]);

                // F_i -> F_(i+1)
                net.addArc("F" + i, "F" + (i + 1));
            }
            // G_i -> H_i
            net.addArc(nodes[nodeIx], nodes[nodeIx + 1]);

            // G -> F
            net.addArc(nodes[nodeIx], finishNode);
            // H -> F
            net.addArc(nodes[nodeIx + 1], finishNode);
            // F -> G
            net.addTemporalArc(finishNode, nodes[nodeIx], 1);
            // F -> H
            net.addTemporalArc(finishNode, nodes[nodeIx + 1], 1);

            // graphisc
            // TODO net.setNodePosition(finishNode, x, y, width, height);
        }
        // observation O
        addNode(net, nodes[nodes.length - 1], levelActions.get(levelActions.size() - 1));
        // H_last -> O
        net.addArc(nodes[nodes.length - 2], nodes[nodes.length - 1]);

        // add temporal arcs
        adjustNodes(net, nodes);

        return net;

    }

    static void adjustNodes(Network net, String[] nodes) {
        // values for layout
        int x = 180;
        int y = 60;
        int width = 80;
        int height = 60;
        for (String node : nodes) {
            //net.setNodeTemporalType(node, Network.NodeTemporalType.Plate);
            if(!node.equals("O")) net.addTemporalArc(node, node, 1);

            // delete default outcomes
            deleteDefaultOutcomes(net, node);

            //net.deleteOutcome(node, 0);
            //net.deleteOutcome(node, 1);

            // do layouting
            net.setNodePosition(node, x, y, width, height);
            y += 80;
        }
    }

    public static void addNode(Network net, String name, int statesNum) {
        addTemporalNode(net, name);
        for (int i = 0; i < statesNum; i++) {
            net.addOutcome(name, "HiddenState" + i);
        }
    }

    public static String[] createNodeNames(int levels) {
        String[] nodes = new String[levels * 2 - 1];
        for (int i = 0; i < levels - 1; i++) {
            int id = levels - i - 1;
            int nodeIx = 2 * i;
            // goal node
            nodes[nodeIx] = "G" + id;
            // hidden node representing goal's internal structure
            nodes[nodeIx + 1] = "H" + id;

        }
        nodes[nodes.length - 1] = "O";
        return nodes;
    }
}
